perf(deepseek_v4): faster exact DeepSeek-V4 prefill (wsdpa fusion + scheduler/cache amortization) - #2568
Conversation
|
Follow-up validation / fix in
Current cold-server measurements on this machine with the native kernel path, WSDPA top-k enabled, 2,048-token blocks, and |
|
Follow-up validation after the native-kernel build fix:
Validation: native MoE parity tests passed; all five WSDPA/trimmed-rotating-cache regression tests passed; a live temperature-0 arithmetic response remained correct. The older ~470 tok/s values in the original draft could not be reproduced under this clean, cache-safe contract, so please use the measured 427.1 tok/s high-context figure for release/PR notes. The PR remains open with the build fix, the long-context gibberish fix, and this verified incremental MoE improvement. |
jonathan308
changed the title
|
Validation complete — reopening. Summary of what changed since this was closed: What was validated
Honest benchmark correction
The upstream cap comes from a measured ~1.6 s fixed cost per 512-token rotating-block boundary (snapshot + stream sync + cache clear + scheduler step); the 2048 alignment in Everything in the PR body is updated to match. Happy to run any additional validation you'd like to see. |
|
I traced the three failing CI jobs to the same test-fixture omission and opened a one-line patch against this PR's head branch: jonathan308#1 The patch adds Verification on the exact patch:
The original CI run is the negative control: the same test fails identically on Python 3.11, 3.12, and 3.13 after 7,982 other tests pass. |
|
CI fix merged: @PhilipJohnBasile correctly traced the 3 failing jobs to a test-fixture omission — |
|
@jonathan308, Thank you for all the work on this one. The wsdpa kernel in particular turned out to be the fastest DeepSeek V4 attention path I have measured so far, so I would really like to get the good parts of this PR in. Since it was opened, #2559, #2562, #2563 and a follow-up (cfec502, DeepSeek V4 cache blocks aligned to 2048) landed on main, which changes the picture a bit, so here is where each commit stands. adbf4f5: main now covers this through cfec502, so only the setup.py CMake pinning part is still needed, and that part is a nice fix. One note from my machine: the ~1.6 s per 512-boundary cost did not reproduce here, the win seems to come from chunk size itself. 981f1d6: I think this one can go. #2562 tunes the same constant, and at 2048-token chunks both versions pick BM16 anyway, so it would not add anything on top. 95b27cb + e3c6912 + f601239: these three are the keepers. I cherry-picked them onto current main and measured (M3 Ultra, V4-Flash oQ4e, cache on, cold runs, tg128):
tg unchanged. Your JIT kernel beats the precompiled deepseek_v4_sparse_attention head to head, so the dispatch should try wsdpa first, and it also helps installs without the compiled extension. Two small asks for the rebase: a standard-mask guard like the one #2559 threads through (the kernel rebuilds causality from offsets, so custom-mask callers need to stay on the reference path), and dropping the env-gated bn=128 variant. a50836a: this one is close, it just needs a compatibility gate. fast.py currently passes the new kwargs unconditionally, so an older built _ext raises TypeError: GLM-5.2's native path stops working and the V4 indexer quietly falls back while the startup probe still reports the kernels as available. A signature probe in fast.py plus an abi-probe test update would cover it. The mask math itself checked out nicely. One correction on the claims, which I want to frame carefully: temp-0 outputs do change on top of main's native ratio-128 path, so tolerance-level parity is the accurate description rather than exact. To be clear, that is not a quality loss. Only the floating-point reduction order differs, greedy flips only happen where the model is effectively indifferent between the top tokens, and the error is far below the checkpoint's own quantization noise. The already-merged native path landed under exactly the same standard. It is just worth a release-note line, since identical prompts will produce different outputs after the upgrade. I know this PR has been through several rounds already, and the remaining pieces are worth the effort. Happy to rerun the same A/B on your rebased branch whenever it is ready. |
Chunked prefill discards every chunk's logits (the prompt's final token is scored by the first decode step), so the full-vocabulary projection per chunk is pure waste for the 129k-vocab DeepSeek-V4 head. Detect patched models that accept skip_lm_head=True and skip the GEMM for intermediate prefill chunks while preserving DSpark and prompt-priming capture side effects. Lossless: temp-0 outputs match the pre-change build exactly at 10k/40k/75k-token agentic prompts. Unknown models keep stock behavior via signature detection, covered by new regression tests.
PoolingCache grew by full-history mx.concatenate at every prefill chunk, copying the entire pooled prefix per chunk (~1GB/chunk at 100k context, ~5GB/chunk at 500k across 62 caches) plus allocator churn. Replace with a capacity-grown backing buffer and slice-assign appends; snapshot/delta extraction still materializes contiguous copies immune to later appends. Bit-identical: temp-0 outputs match the pre-change build exactly at 10k/40k/75k-token agentic prompts, including restore of SSD cache blocks written by the previous implementation. Adds a startup probe logging native indexer-kernel availability. Cold-prefill throughput +1-3% with the largest effect at 100k+ context.
…ernel The indexer call site applied the PoolingCache causal mask via a separate mx.where pass over the full [B,1,L,P] scores tensor (~5.3GB of extra DRAM traffic per chunk at 100k context). Teach dsa_indexer_scores optional mask_ratio/mask_q_offset kwargs so the kernel epilogue writes the same finfo.min sentinel itself, and use it when the mask is the plain 2-D PoolingCache ratio mask. Batched 3-D masks and the no-mask case keep the previous behavior. Compatibility: fast.py probes the built extension's dsa_indexer_scores signature once at import (_probe_mask_fold). Builds predating the fold kwargs keep the historical call signature — so GLM-5.2's native path is unaffected — and the same mask is applied in a second pass with identical sentinel semantics. Covered by abi-probe tests for old/new/missing signatures and exact second-pass masking on pre-fold builds. Bit-identical: same fp32 MMA K-order, ReLU-before-weight, ascending-h fp32 head-sum, single RNE bf16 cast, identical sentinel; verified by fused-vs-unfused bitwise regression tests (bf16/fp16, ratio 4/128, offset 0/>0) and temp-0 output equivalence at 10k/40k/75k-token agentic prompts. The indexer tile config is fixed at bn=64/wm2/wn2: a bn=128 variant measured no win (compute/barrier-bound) and is removed.
80bead1 to
9423196
Compare
|
@jundot Rebased onto current main with every ask addressed — ready for your A/B rerun whenever you'd like. Per-commit:
Focused suite on the rebased branch: 618 passed, including the fused-mask bit-identity tests on a freshly built extension and the trimmed-cache coordinate regression tests. |
|
Reran the A/B on the rebased branch (M3 Ultra, V4-Flash oQ4e, cache on, cold runs, tg128):
tg unchanged, and the numbers land right where the wsdpa-first slice predicted. I also verified the compatibility gate against a real pre-fold extension build: the probe detects it, the historical call signature is kept so GLM-5.2 stays on its native path, and the second-pass mask is bit-identical to the fused semantics. Greedy outputs shift as expected from the kernel swap and are deterministic per build. Everything I asked for is in, so this is good to merge on my side. Merging now, and I will pick up a couple of cosmetic leftovers in a follow-up (the unused bn=128 metal instantiations and an inline import). |
3 participants
Summary
Exact-preserving DeepSeek-V4 prefill speedup, rebased onto current main (post-#2559/#2562/#2563/cfec5021). Six commits:
c2f55836fuse windowed+pooled prefill attention — JIT Metal kernel for the exact DSV4 shape (B=1, H=64, D=512, bf16), routed first forLocalAttention,CompressedAttention, andSparseCompressedAttention(small-pool and indexer top-k). Dispatch is gated on a_standard_maskflag threaded from the model loop (same pattern as perf(deepseek-v4): route ratio-128 prefill through native attention #2559): the kernels rebuild causality/sliding-window/pooled visibility from offsets, so direct callers with custom masks stay on the reference path.d372749dsetup.py CMake pinning so the native prefill kernels build reliably (the block-alignment half of the original commit is covered by cfec502 and was dropped).19827b93skip discardedlm_headprojection in intermediate chunked-prefill chunks (logits of non-final chunks are thrown away).a77ab09eappend pooled cache in place during prefill (removes per-chunk pooled-prefix re-concatenation).760fe945fold the pooled-ratio causal mask into the indexer score kernel epilogue, with a signature probe (_probe_mask_fold): extension builds predating the fold kwargs keep the historical call signature (GLM-5.2 unaffected) and get the same mask in a second pass with identical sentinel semantics. The indexer tile config is fixed at bn=64; the env-gated bn=128 variant measured no win and was removed.9423196atest fixture compatibility for theskip_lm_headcapability check (thanks @PhilipJohnBasile).The MXFP4 tile-tuning commit from the previous revision was dropped: #2562 tunes the same constant and both pick BM16 at 2048-token chunks.
No precision, selection-quality, sampling, or cache-semantics changes. Decode, DSpark verify, batching, cache update, prompt-boundary snapshots, and fallback behavior are unchanged; every kernel path falls back to stock SDPA/native on shape-gate miss, setup failure, or runtime error. Escape hatches:
OMLX_DSV4_WSDPA=0,OMLX_DSV4_WSDPA_TOPK=0.Performance
@jundot's A/B on current main (M3 Ultra, V4-Flash oQ4e, cache on, cold runs, tg128), cherry-picking the three core commits:
tg unchanged. The JIT kernel also beats the precompiled
deepseek_v4_sparse_attentionhead to head and helps installs without the compiled extension.The remaining high-context taper is architectural: the DSA indexer scores all context/4 pooled rows per query, so prefill FLOPs grow linearly with absolute context. Kernel-traffic variants (K-tile hoist, bn=128) were implemented and measured at 0% gain — excluded.
Output-parity note (release-note worthy)
Greedy/temp-0 outputs after this PR are tolerance-level parity, not bit-identical, with current main's native ratio-128 path: floating-point reduction order differs, so greedy flips can occur where the model is effectively indifferent between top tokens — far below the checkpoint's own quantization noise (the same standard the merged native path landed under). Each change was validated temp-0 exact-match against its own pre-change build at 10k/40k/75k-token prompt sets, and the fused kernels are bit-exact against explicit fp32 references (including trimmed-RotatingKVCache coordinate regression tests — both fail pre-fix, pass post-fix).
Upgrade note for the release log: identical prompts can produce different (equivalent-quality) outputs after this change.
Validation
test_deepseek_v4_wsdpa,test_dsa_indexer_fused_mask(bit-identity incl. fresh bn=64 build),test_custom_kernel_abi_probe(old/new/missing mask-fold signatures + exact second-pass masking),test_pooling_cache_*,test_scheduler*,test_deepseek_v4_patch,test_deepseek_v4_dspark,test_glm_*,test_prefill_oom_graceful)Evaluated and deliberately excluded (with data): indexer 64-alignment padding (~0.5% gain, shifts temp-0 numerics),
mx.compileMoE wrapper (0%),mx.compiledecode step (unvalidated state risk), K-tile/bn=128 indexer variants (0%, FLOP-bound), MXFP4 tile re-tune (superseded by #2562).